import torch
from image_process.EyeDataset import EyeDataset
from image_process.process import image_process
from image_process.ROI import try_roi_process
from image_process.noise_manage import events_to_count_image_single
from network.calibration import calibrate_eye_position, apply_offset_to_events
import numpy as np

def batch_ellipse_to_features_torch(ellipse_vecs: torch.Tensor):
    """
    批量将椭圆参数向量转换为标准化的21维特征向量
    :param ellipse_vecs: shape (N, 5) torch.Tensor
    :return: shape (N, 21) 标准化后的特征 torch.Tensor
    """
    N = ellipse_vecs.shape[0]
    aug_vecs = torch.cat([ellipse_vecs, torch.ones(N, 1, device=ellipse_vecs.device)], dim=1)  # (N, 6)
    matrices = aug_vecs.unsqueeze(2) * aug_vecs.unsqueeze(1)  # (N, 6, 6)
    idx = torch.tril_indices(6, 6, device=ellipse_vecs.device)
    features = matrices[:, idx[0], idx[1]]  # (N, 21)
    mean = features.mean(dim=0, keepdim=True)
    std = features.std(dim=0, keepdim=True)
    std[std == 0] = 1
    features_std = (features - mean) / std
    return features_std

class GazeDataset(EyeDataset):
    def __init__(self, data_dir, user, seq_len=2000, delta=10, min_score=0.5, img_shape=(240, 346)):
        """
        :param data_dir: 数据目录
        :param user: 用户编号
        :param seq_len: 每组事件数量
        :param delta: 滑动窗口步长
        :param min_score: 椭圆拟合最小分数
        :param img_shape: 图像尺寸
        """
        super().__init__(data_dir, user)
        self.collect_data()
        self.seq_len = seq_len
        self.delta = delta
        self.min_score = min_score
        self.frame_list = self.frame_stack[::-1]  # 时间从旧到新
        self.event_list = self.event_stack[::-1]  # 时间从旧到新

        # ----------- 初始校准 -----------
        # 取前4000个事件做位置校准
        init_events = self.event_list[:4000]
        offset_x, offset_y = calibrate_eye_position(init_events, img_shape=img_shape)
        # 对所有事件应用偏移
        self.event_list = apply_offset_to_events(self.event_list, offset_x, offset_y, img_shape=img_shape)
        # --------------------------------

        self.ellipse_list = []
        self.label_list = []

        event_idx = 0
        frame_idx = 0
        num_events = len(self.event_list)
        num_frames = len(self.frame_list)

        while event_idx + seq_len <= num_events:
            # 1. 取当前窗口的事件
            event_window = self.event_list[event_idx:event_idx+seq_len]

            # 2. 椭圆拟合
            roi_result = try_roi_process(*events_to_count_image_single(event_window), ellipse=self.ellipse_list[:-1])
            if roi_result is not None:
                ellipse = roi_result
            else:
                ellipse = image_process(event_window, self.min_score)
                
            if ellipse is not None:
                # 3. 找到该区间内最新帧的标签
                last_event_ts = event_window[-1]
                label = None
                for i in range(frame_idx, num_frames):
                    frame = self.frame_list[i]
                    if frame.timestamp <= last_event_ts:
                        label = (frame.row, frame.col)
                        frame_idx = i
                    else:
                        break
                if label is not None:
                    self.ellipse_list.append(ellipse)
                    self.label_list.append(label)
            # 4. 滑动窗口
            event_idx += delta

    def __len__(self):
        return len(self.ellipse_list)

    def __getitem__(self, idx):
        '''
        获取指定索引的特征和标签
        :return: 返回标准化的21维特征向量和标签
        '''
        ellipse = torch.tensor(self.ellipse_list[idx], dtype=torch.float32)
        features = batch_ellipse_to_features_torch(ellipse.unsqueeze(0)).squeeze(0)
        label = torch.tensor(self.label_list[idx], dtype=torch.float32)
        return features, label

def split_dataset(dataset, train_ratio=0.7, val_ratio=0.15, test_ratio=0.15, seed=42):
    """
    按照给定比例划分数据集为训练集、验证集和测试集
    :param dataset: GazeDataset实例
    :param train_ratio: 训练集比例
    :param val_ratio: 验证集比例
    :param test_ratio: 测试集比例
    :param seed: 随机种子
    :return: (train_idx, val_idx, test_idx)
    """
    np.random.seed(seed)
    total_len = len(dataset)
    indices = np.arange(total_len)
    np.random.shuffle(indices)
    train_end = int(total_len * train_ratio)
    val_end = train_end + int(total_len * val_ratio)
    train_idx = indices[:train_end]
    val_idx = indices[train_end:val_end]
    test_idx = indices[val_end:]
    return train_idx, val_idx, test_idx
